import os
import shutil

def ensure_weights():
    here = os.path.dirname(os.path.abspath(__file__))
    dst_dir = os.path.join(here, 'models', 'weights')
    os.makedirs(dst_dir, exist_ok=True)
    dst = os.path.join(dst_dir, 'model_final.pth')

    if os.path.isfile(dst):
        print(f"[ok] 权重已存在: {dst}")
        return

    # 尝试从上级项目目录或用户指定路径拷贝
    candidates = []
    env_src = os.environ.get('WEIGHTS_SRC')
    if env_src:
        candidates.append(env_src)
    candidates += [
        os.path.join(here, '..', 'weights', 'model_final.pth'),
        os.path.join(here, '..', 'models', 'weights', 'model_final.pth'),
    ]

    for src in candidates:
        if src and os.path.isfile(src):
            print(f"[copy] 从 {src} 拷贝到 {dst}")
            shutil.copy2(src, dst)
            print("[ok] 权重拷贝完成")
            return

    print("[hint] 未找到权重。请将 model_final.pth 放入 deploy/models/weights/，或设置 WEIGHTS_SRC 指向权重路径。")

if __name__ == '__main__':
    ensure_weights()